#!/usr/bin/python3
'''
Contributor: Catfood
Date: 2021.3.7 ~ 2021.3.14
'''

'''
What f***
now we have 4 cameras
transform matrix order:
0 -> U <-> cam U
1 -> R <-> cam U
2 -> L <-> cam D
3 -> D <-> cam D
4 -> F <-> cam F
5 -> B <-> cam B
A little bit different from the legacy script
we start from 0 in operation, consistent with matrix order
we use key 0-5 to save perspective mat
we use key a-f to save perspectived picture
'''

import numpy as np
import cv2
import configparser
import pickle
import sys, time

import Camera
import Pretreat

current_point_counter = -1
colors = [(0, 0, 255), (0, 255, 0), (255, 255, 0), (255, 0, 0)]
radius = [0, 0, 0, 0]
points = []
frame = None
M = None

def test_pretreat():
    cam_idx_list, calib_mtx, calib_dist = Camera.load_camera_params()

    config = configparser.ConfigParser()
    config.read('../configs/vision_pretreat.ini')
    treator = Pretreat.Pretreat(config)

    #caps = [cv2.VideoCapture(cam_idx_list[i]) for i in range(len(cam_idx_list))]
    #for cap in caps:
    #    cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M','J','P','G'))
        # cap.set(cv2.CAP_PROP_FRAME_WIDTH, Camera.WIDTH)
        # cap.set(cv2.CAP_PROP_FRAME_HEIGHT, Camera.HEIGHT)
#    for i in range(len(cam_idx_list)):
        # caps[i].release()
#        caps[i].open(cam_idx_list[i])
#         caps[i].read()
    caps = []
    #cam_idx_list.reverse()
    # cam_idx_list = [6, 0, 4, 2]
    for i in range(len(cam_idx_list)):
        caps.append(Camera.get_capture())
        caps[i].open(cam_idx_list[i])
    while True:
        stat = [cap.isOpened() for cap in caps]
        if False not in stat:
            break
    print("camera ok")

    while True:
        # to ensure that the camera is reading images
        time.sleep(0.5)
        frames = []

        for cp_open in range(len(cam_idx_list)):
            print(cp_open, cam_idx_list[cp_open])
            frames.append(caps[cp_open].read()[1])
            cv2.imshow(str(cp_open), frames[cp_open])
            time.sleep(0.5)
#           frames[cp_open] = Camera.camera_calibration(frames[cp_open], calib_mtx[cp_open], calib_dist[cp_open])

        treator.raw_four_images = frames # black magic; Just for test
        treator.CutImage(test=True)
        perspectived = treator.perspectived_imgs
        huge_image = np.zeros((3 * 300, 4 * 300, 3)).astype('uint8')
        huge_image[1 * 300 : 2 * 300, 0 * 300 : 1 * 300, :] = perspectived[5 - 1]
        huge_image[0 * 300 : 1 * 300, 1 * 300 : 2 * 300, :] = perspectived[1 - 1]
        huge_image[1 * 300 : 2 * 300, 1 * 300 : 2 * 300, :] = perspectived[3 - 1]
        huge_image[2 * 300 : 3 * 300, 1 * 300 : 2 * 300, :] = perspectived[4 - 1]
        huge_image[1 * 300 : 2 * 300, 2 * 300 : 3 * 300, :] = perspectived[2 - 1]
        huge_image[1 * 300 : 2 * 300, 3 * 300 : 4 * 300, :] = perspectived[6 - 1]

        cv2.imshow("Fin", huge_image)
        key = cv2.waitKey(1)
        if key == ord('q'):
            break
    cv2.destroyAllWindows()

def MouseHandler(event, x, y, flags, param):
    global current_point_counter
    global points
    if event == cv2.EVENT_LBUTTONDOWN :
        current_point_counter += 1
        current_point_counter %= 4
        if len(points) < 4:
            points.append((x, y))
        else:
            points[current_point_counter] = (x, y)
            radius[current_point_counter] = 0
    return 

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == 'test':
        test_pretreat()
        exit()

    CP_OPEN = 0
    #cp = cv2.VideoCapture(0)
    
    cam_idx_list, calib_mtx, calib_dist = Camera.load_camera_params()

    cp = Camera.get_capture()
    cv2.namedWindow("233")
    cv2.setMouseCallback("233", MouseHandler)
    print(cam_idx_list)
    while True:
        cp.open(cam_idx_list[CP_OPEN])
        _, frame = cp.read()
        #frame = Camera.camera_calibration(frame, calib_mtx[CP_OPEN], calib_dist[CP_OPEN])
        cv2.imshow("233", frame)
        
        for i in range(current_point_counter+1):
            if radius[i] < 7:
                radius[i] += 1
            cv2.circle(frame, points[i], radius[i], colors[i], 2)
            frame = cv2.putText(frame, str(i+1), points[i], cv2.FONT_HERSHEY_SIMPLEX, 0.77, (255, 255, 255), 2)
            cv2.imshow("233", frame)
        
        if current_point_counter == 3:
            pts1 = np.float32(points)
            pts2 = np.float32([[0,0],[300,0],[0,300],[300,300]]) 
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   
            M = cv2.getPerspectiveTransform(pts1,pts2)
            dst = cv2.warpPerspective(frame,M,(300,300))

            cv2.imshow("fixed", dst)
        key = cv2.waitKey(1)
        if key == ord('q'):
            break
        elif key == ord('c'):
            cp.release()
            CP_OPEN += 1
            CP_OPEN %= 4
            print(CP_OPEN)
            cp.open(cam_idx_list[CP_OPEN])
        elif key-ord('0') <= 5 and key != -1:
            config = configparser.ConfigParser()
            config.read("../configs/vision_pretreat.ini")
            print(config.sections())
            section = 'tran_mat_%d'%(key - ord('0'))
            if section not in config.sections():
                config.add_section(section)
            for i in range(9):
                config[section]['element_%d'%i] = str(M[int(i/3), int(i%3)])
            with open("../configs/vision_pretreat.ini", "w") as r:
                config.write(r)
        # collect images
        elif key-ord('a') <= 5 and key-ord('a') >= 0:
            cv2.imwrite("../configs/%d.jpg"%(key-ord('a')), frame)
